//
// Created by HP on 2022/1/19.
//

#include "lve_model.h"
#include "lve_utils.h"

// libs
#define TINYOBJLOADER_IMPLEMENTATION
#include <tiny_obj_loader.h>
#define GLM_ENABLE_EXPERIMENTAL
#include <glm/gtx//hash.hpp>

// std
#include <unordered_map>

namespace std {
    template <>
    struct hash<lve::LveModel::Vertex> {
        size_t operator() (lve::LveModel::Vertex const &vertex) const {
            size_t seed = 0;
            lve::hashCombine(seed, vertex.position, vertex.color, vertex.normal, vertex.uv);
            return seed;
        }
    };
}

namespace lve {

    LveModel::LveModel(LveDevice &device, const Builder &builder): lveDevice{device} {
        createVertexBuffers(builder.vertices);
        createIndexBuffers(builder.indices);
    }

    LveModel::~LveModel() {
        // pAllocator: 为内存分配指定分配器
        // 原因: 在渲染高复杂的场景时, 会有数千个实例等待渲染, 此时要进行大量的内存分配操作
        // GPU 在某一时间分配内存的次数有上限, 使用 allocator 可自己定义内存分配行为, 如:
        // allocator 向系统一次性分配大块内存, 再通过内存池等方式分割、管理内存, 则可避免此中情况
        vkDestroyBuffer(lveDevice.device(), vertexBuffer, nullptr);
        vkFreeMemory(lveDevice.device(), vertexBufferMemory, nullptr);
        if (hasIndexBuffer) {
            vkDestroyBuffer(lveDevice.device(), indexBuffer, nullptr);
            vkFreeMemory(lveDevice.device(), indexBufferMemory, nullptr);
        }
    }

    std::unique_ptr<LveModel> LveModel::createModelFromFile(LveDevice &device, const std::string &filepath) {
        Builder builder{};
        builder.loadModel(filepath);
        return std::make_unique<LveModel>(device, builder);
    }

    void LveModel::createVertexBuffers(const std::vector<Vertex> &vertices) {
        vertexCount = static_cast<uint32_t>(vertices.size());
        assert(vertexCount >= 3 && "Vertex count must be at least 3.");
        VkDeviceSize bufferSize = sizeof(vertices[0]) * vertexCount;

        VkBuffer stagingBuffer{};
        VkDeviceMemory stagingBufferMemory{};
        lveDevice.createBuffer(
            bufferSize,
            VK_BUFFER_USAGE_TRANSFER_SRC_BIT, // 指明这个 Buffer 是作为传输内存的 src buffer
            // Host: CPU, Device: GPU
            // HOST_VISIBLE: 申请可以在 Host 上访问的内存, HOST_COHERENT: 使得 Host 和 Device 上的内存数据保存一致
            VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,
            stagingBuffer,
            stagingBufferMemory
        );
        // 为显存和内存作映射
        void *data;
        vkMapMemory(lveDevice.device(), stagingBufferMemory, 0, bufferSize, 0, &data);
        // 写入数据
        // stagingBuffer 位于显卡的一个类似于缓冲区的地方
        memcpy(data, vertices.data(), static_cast<size_t>(bufferSize));
        // 解除映射
        vkUnmapMemory(lveDevice.device(), stagingBufferMemory);
        lveDevice.createBuffer(
            bufferSize,
            // 指明这个 Buffer 是作为传输内存的 des     | 指明这个 Buffer 是用了储存 Vertex 的
            VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_VERTEX_BUFFER_BIT,
            // Host: CPU, Device: GPU
            // DEVICE_LOCAL_BIT: 显卡的本地内存
            VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT,
            vertexBuffer,
            vertexBufferMemory
        );
        // 将位于缓冲区的 stagingBuffer 的数据拷贝到显存中
        lveDevice.copyBuffer(stagingBuffer, vertexBuffer, bufferSize);

        // 清理 stagingBuffer
        vkDestroyBuffer(lveDevice.device(), stagingBuffer, nullptr);
        vkFreeMemory(lveDevice.device(), stagingBufferMemory, nullptr);
    }

    void LveModel::createIndexBuffers(const std::vector<uint32_t> &indices) {
        indexCount = static_cast<uint32_t>(indices.size());
        hasIndexBuffer = indexCount > 0;
        if (!hasIndexBuffer) {
            return;
        }
        VkDeviceSize bufferSize = sizeof(indices[0]) * indexCount;

        VkBuffer stagingBuffer{};
        VkDeviceMemory stagingBufferMemory{};
        lveDevice.createBuffer(
            bufferSize,
            VK_BUFFER_USAGE_TRANSFER_SRC_BIT, // 指明这个 Buffer 是作为传输内存的 src buffer
            // Host: CPU, Device: GPU
            // HOST_VISIBLE: 申请可以在 Host 上访问的内存, HOST_COHERENT: 使得 Host 和 Device 上的内存数据保存一致
            VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,
            stagingBuffer,
            stagingBufferMemory
        );
        // 为显存和内存作映射
        void *data;
        vkMapMemory(lveDevice.device(), stagingBufferMemory, 0, bufferSize, 0, &data);
        // 写入数据
        // stagingBuffer 位于显卡的一个类似于缓冲区的地方
        memcpy(data, indices.data(), static_cast<size_t>(bufferSize));
        // 解除映射
        vkUnmapMemory(lveDevice.device(), stagingBufferMemory);
        lveDevice.createBuffer(
                bufferSize,
                // 指明这个 Buffer 是作为传输内存的 des     | 指明这个 Buffer 是用了储存 index 的
                VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_INDEX_BUFFER_BIT,
                // Host: CPU, Device: GPU
                // DEVICE_LOCAL_BIT: 显卡的本地内存
                VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT,
                indexBuffer,
                indexBufferMemory
                );
        // 将位于缓冲区的 stagingBuffer 的数据拷贝到显存中
        lveDevice.copyBuffer(stagingBuffer, indexBuffer, bufferSize);

        // 清理 stagingBuffer
        vkDestroyBuffer(lveDevice.device(), stagingBuffer, nullptr);
        vkFreeMemory(lveDevice.device(), stagingBufferMemory, nullptr);
    }

    void LveModel::bind(VkCommandBuffer commandBuffer) {
        VkBuffer buffers[] = {vertexBuffer};
        VkDeviceSize offsets[] = {0};
        vkCmdBindVertexBuffers(commandBuffer, 0, 1, buffers, offsets);

        if (hasIndexBuffer) {
            vkCmdBindIndexBuffer(commandBuffer, indexBuffer, 0, VK_INDEX_TYPE_UINT32);
        }
    }

    void LveModel::draw(VkCommandBuffer commandBuffer) const {
        if (hasIndexBuffer) {
            vkCmdDrawIndexed(commandBuffer, indexCount, 1, 0, 0, 0);
        } else {
            vkCmdDraw(commandBuffer, vertexCount, 1, 0, 0);
        }
    }

    std::vector<VkVertexInputBindingDescription> LveModel::Vertex::getBindDescription() {
        std::vector<VkVertexInputBindingDescription> bindingDescriptions(1);
        bindingDescriptions[0].binding = 0;
        bindingDescriptions[0].stride = sizeof(Vertex);
        bindingDescriptions[0].inputRate = VK_VERTEX_INPUT_RATE_VERTEX;
        return bindingDescriptions;
    }

    std::vector<VkVertexInputAttributeDescription> LveModel::Vertex::getAttributeDescriptions() {
        std::vector<VkVertexInputAttributeDescription> attributeDescriptions(3);
        attributeDescriptions[0].binding = 0;
        attributeDescriptions[0].location = 0;
        attributeDescriptions[0].format = VK_FORMAT_R32G32B32_SFLOAT;
        attributeDescriptions[0].offset = offsetof(Vertex, position);

        attributeDescriptions[1].binding = 0;
        attributeDescriptions[1].location = 1;
        attributeDescriptions[1].format = VK_FORMAT_R32G32B32_SFLOAT;
        attributeDescriptions[1].offset = offsetof(Vertex, color);

        attributeDescriptions[2].binding = 0;
        attributeDescriptions[2].location = 2;
        attributeDescriptions[2].format = VK_FORMAT_R32G32B32_SFLOAT;
        attributeDescriptions[2].offset = offsetof(Vertex, normal);

        attributeDescriptions[3].binding = 0;
        attributeDescriptions[3].location = 3;
        attributeDescriptions[3].format = VK_FORMAT_R32G32_SFLOAT;
        attributeDescriptions[3].offset = offsetof(Vertex, uv);

        return attributeDescriptions;
    }

    void LveModel::Builder::loadModel(const std::string &filepath) {
        tinyobj::attrib_t attrib{}; // 储存 顶点坐标、颜色、法线、纹理坐标
        std::vector<tinyobj::shape_t> shapes{}; // 储存 index
        std::vector<tinyobj::material_t> materials{}; // 材质
        std::string warn, err;

        if (!tinyobj::LoadObj(&attrib, &shapes, &materials, &warn, &err, filepath.c_str())) {
            throw std::runtime_error(warn + err);
        }

        vertices.clear();
        indices.clear();

        std::unordered_map<Vertex, uint32_t> uniqueVertices{};
        for (const auto& shape: shapes) {
            for (const auto& index: shape.mesh.indices) {
                Vertex vertex{};
                // 因为 index 是可选的
                // 所以使用 -1 表示没有提供该数据
                if (index.vertex_index >= 0) {
                    vertex.position = {
                        attrib.vertices[3 * index.vertex_index + 0],
                        attrib.vertices[3 * index.vertex_index + 1],
                        attrib.vertices[3 * index.vertex_index + 2],
                    };

                    vertex.color = {
                        attrib.colors[3 * index.vertex_index + 0],
                        attrib.colors[3 * index.vertex_index + 1],
                        attrib.colors[3 * index.vertex_index + 2],
                    };
                }
                if (index.normal_index >= 0) {
                    vertex.normal = {
                        attrib.normals[3 * index.normal_index + 0],
                        attrib.normals[3 * index.normal_index + 1],
                        attrib.normals[3 * index.normal_index + 2],
                    };
                }
                if (index.texcoord_index >= 0) {
                    vertex.uv = {
                        attrib.texcoords[2 * index.texcoord_index + 0],
                        attrib.texcoords[2 * index.texcoord_index + 1],
                    };
                }

                if (uniqueVertices.count(vertex) == 0) {
                    uniqueVertices[vertex] = static_cast<uint32_t>(vertices.size());
                    vertices.push_back(vertex);
                }
                indices.push_back(uniqueVertices[vertex]);
            }
        }
    }
}